package padding

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"encoding/base64"
	"fmt"
)

const (
	PKCS5PaddingAESKeyLength = 16
	ZeroPaddingAesKeyLength  = 16
)

func PKCS5PaddingEncode(data, aesKey string) (string, error) {
	if aesKey == "" {
		return data, nil
	}
	if len([]byte(aesKey)) != PKCS5PaddingAESKeyLength {
		aesKey = fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[8:24]
	}
	key := []byte(aesKey)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	iv := []byte(fmt.Sprintf("%x", md5.Sum(key))[:16])
	mode := cipher.NewCBCEncrypter(block, iv)
	padData := pkcs5Padding([]byte(data), block.BlockSize())
	cipherData := make([]byte, len(padData))
	mode.CryptBlocks(cipherData, padData)
	return base64.StdEncoding.EncodeToString(cipherData), nil
}

func PKCS5PaddingDecode(data, aesKey string) (string, error) {
	if aesKey == "" {
		return data, nil
	}
	if len([]byte(aesKey)) != PKCS5PaddingAESKeyLength {
		aesKey = fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[8:24]
	}
	key := []byte(aesKey)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	iv := []byte(fmt.Sprintf("%x", md5.Sum(key))[:16])
	mode := cipher.NewCBCDecrypter(block, iv)
	decodeData, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}
	plainData := make([]byte, len(decodeData))
	mode.CryptBlocks(plainData, decodeData)
	unpadData := pkcs5Unpadding(plainData)
	return string(unpadData), nil
}

func pkcs5Padding(data []byte, blockSize int) []byte {
	padding := blockSize - (len(data) % blockSize)
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

func pkcs5Unpadding(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:(length - unpadding)]
}

func ZeroPaddingEncode(data string, aesKey string) (string, error) {
	if aesKey == "" {
		return data, nil
	}
	if len(aesKey) != ZeroPaddingAesKeyLength {
		aesKey = fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[8:24]
	}
	ivStr := fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[0:16]
	iv := []byte(ivStr)
	key := []byte(aesKey)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	dataBytes := []byte(data)
	paddedData := zeroPad(dataBytes, block.BlockSize())
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(paddedData, paddedData)
	//return safeBase64Encode(paddedData), nil
	rs := base64.StdEncoding.EncodeToString(paddedData)
	return rs, nil
}

func ZeroPaddingDecode(data string, aesKey string) (string, error) {
	if aesKey == "" {
		return data, nil
	}
	if len(aesKey) != ZeroPaddingAesKeyLength {
		aesKey = fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[8:24]
	}
	ivStr := fmt.Sprintf("%x", md5.Sum([]byte(aesKey)))[0:16]
	iv := []byte(ivStr)
	key := []byte(aesKey)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	ciphertext, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)
	unpaddedData := zeroUnpad(ciphertext)
	return string(unpaddedData), nil
}

// Zero padding functions
func zeroPad(data []byte, blockSize int) []byte {
	padding := blockSize - (len(data) % blockSize)
	padText := bytes.Repeat([]byte{byte(0)}, padding)
	return append(data, padText...)
}

func zeroUnpad(data []byte) []byte {
	for i := len(data) - 1; i >= 0; i-- {
		if data[i] != 0 {
			return data[:i+1]
		}
	}
	return nil
}
